#ifndef DEFECTRECOGNIZER_H
#define DEFECTRECOGNIZER_H

#include<opencv2/opencv.hpp>
#include<QString>
#include "yolo.h"

class DefectRecognizer
{
public:
    DefectRecognizer();

    // 识别缺陷的分类器
private:
    Yolov8                      m_objYolov8;
    cv::dnn::Net                m_objNet3;
    std::vector<cv::Scalar>     m_ObjColor;
    QString                     m_strYoloOnnxModelPath;
    QString                     m_strYoloOnnxClassesPath;
    bool                        m_bUseGPU;

public:
    void                        SetYoloOnnxModelPath(QString strPath){m_strYoloOnnxModelPath = strPath;}

private:
    int                         m_nImageHeight;
    int                         m_nImageWidth;
    int                         m_nPartHeight;
    int                         m_nPartWidth;

    int                         m_nCountHeight;
    int                         m_nCountWidth;

private:
    void                        AdjustPartSize();

    void                        InitYoloParameter();

public:
    void                        LoadYoloModel();

public:
    void                        SetImageHeight(int value);
    void                        SetImageWidth(int value);
    void                        SetPartHeight(int value);
    void                        SetPartWidth(int value);


    std::vector<Detection>      RecognizePartsDefect(cv::Mat image);
    std::vector<Detection>      RecognizeDefect(cv::Mat image);

    void                        DrawPred(cv::Mat image,std::vector<Detection> results);

    void                        SetUseGPU(bool useGPU);

    void                        test();
};

#endif // DEFECTRECOGNIZER_H
